"""
python p2_solution_particle.py

"""

import json
import os
import numpy as np
from sko.PSO import PSO

# 数据加载路径
warehouse_path = '../fujian/fujian3/origin_data/warehouse.json'
inventory_path = '../fujian/fujian3/data_from_p1/all_average_inventory.json'
sales_path = '../fujian/fujian3/data_from_p1/all_average_sales.json'
output_path = '../fujian/fujian3/particle/'

# 数据加载
with open(warehouse_path, 'r') as f:
    warehouses = json.load(f)

with open(inventory_path, 'r') as f:
    inventory_data = json.load(f)

with open(sales_path, 'r') as f:
    sales_data = json.load(f)

# 创建字典
warehouse_dict = {w['warehouse_id']: w for w in warehouses}
inventory_dict = {i['category_id']: i['average_inventory'] for i in inventory_data}
sales_dict = {s['category_id']: s['average_sales'] for s in sales_data}

categories = list(inventory_dict.keys())
num_categories = len(categories)
warehouses_list = list(warehouse_dict.keys())

# 适应度函数
def fitness_function(solution):
    warehouse_utilization = {w: {'inventory': 0, 'sales': 0} for w in warehouses_list}
    total_cost = 0

    for category_index, warehouse_index in enumerate(solution.astype(int)):
        warehouse_id = warehouses_list[warehouse_index]
        category_id = categories[category_index]
        inventory = inventory_dict[category_id]
        sales = sales_dict[category_id]

        # 更新仓库的库存和出货量
        warehouse_utilization[warehouse_id]['inventory'] += inventory
        warehouse_utilization[warehouse_id]['sales'] += sales

        if inventory > 0:
            total_cost += warehouse_dict[warehouse_id]['daily_cost']

    # 检查约束条件
    for warehouse_id, usage in warehouse_utilization.items():
        if (usage['inventory'] >= warehouse_dict[warehouse_id]['max_inventory'] or
            usage['sales'] >= warehouse_dict[warehouse_id]['max_sales']):
            return float('inf')  # 违反约束，返回无穷大

    return total_cost

# 粒子群优化
def run_pso():
    # 设置 PSO 参数
    # pso = PSO(func=fitness_function, dim=num_categories, n_particles=50, max_iter=100, 
    #           lb=[0]*num_categories, ub=[len(warehouses_list)-1]*num_categories)
    pso = PSO(func=fitness_function, dim=num_categories, population_size=50, max_iter=100, 
          lb=[0]*num_categories, ub=[len(warehouses_list)-1]*num_categories)

    
    best_solution = pso.run()
    return best_solution

# 运行粒子群优化算法
best_solution = run_pso()

# 计算最终结果
total_cost = fitness_function(best_solution)
warehouse_utilization = {w: {'inventory': 0, 'sales': 0} for w in warehouses_list}

for category_index, warehouse_index in enumerate(best_solution.astype(int)):
    warehouse_id = warehouses_list[warehouse_index]
    category_id = categories[category_index]
    inventory = inventory_dict[category_id]
    sales = sales_dict[category_id]
    warehouse_utilization[warehouse_id]['inventory'] += inventory
    warehouse_utilization[warehouse_id]['sales'] += sales

# 输出结果
assignment = [{'category_id': categories[i], 'warehouse_id': warehouses_list[int(best_solution[i])]} for i in range(num_categories)]
os.makedirs(output_path, exist_ok=True)

# 保存分配方案
with open(os.path.join(output_path, 'assignment.json'), 'w') as f:
    json.dump(assignment, f)

# 保存总成本
with open(os.path.join(output_path, 'all_cost.txt'), 'w') as f:
    f.write(str(total_cost))

# 计算和保存利用率
utilization_rates_inventory = []
utilization_rates_sales = []

for warehouse_id, usage in warehouse_utilization.items():
    utilization_rates_inventory.append({
        'warehouse_id': warehouse_id,
        'utilization_rate_of_inventory': usage['inventory'] / warehouse_dict[warehouse_id]['max_inventory']
    })
    utilization_rates_sales.append({
        'warehouse_id': warehouse_id,
        'utilization_rate_of_sales': usage['sales'] / warehouse_dict[warehouse_id]['max_sales']
    })

# 保存库存利用率
with open(os.path.join(output_path, 'all_utilization_rate_inventory.json'), 'w') as f:
    json.dump(utilization_rates_inventory, f)

# 保存出货量利用率
with open(os.path.join(output_path, 'all_utilization_rate_sales.json'), 'w') as f:
    json.dump(utilization_rates_sales, f)

print("完成！最佳总成本:", total_cost)
